#include <tbf.h>

typedef struct {
	int token;
	int cps;
	int burst;
}tbf_t;

//指针数组中的每个指针都指向令牌桶
static tbf_t *jobs[TBF_MAX] = {};
typedef void (*sighandler_t)(int);
//保存alrm信号原本的属性值
sighandler_t alarm_save;
void sig_handler (int s);
static int get_free_pos (void);
void  sig_module_unload (void); 
void  sig_module_load (void);
//为1, 表示已经启动信号模块了 
static int inited = 0;

int tbf_init (int cps, int burst)
{
	tbf_t *me;
	int pos;
	//第一个令牌的初始化 ---> 启动信号模块
	if (!inited) {
		sig_module_load ();
		inited = 1;
	}
	me = malloc (sizeof (tbf_t));
	if (NULL == me) {
		return -ENOMEM;
	}
	me->token = 0;
	me->cps = cps;
	me->burst = burst;
	//找到令牌桶数组中第一个空闲位置 
	pos = get_free_pos ();
	if (-1 == pos) {
		free (me);
		me = NULL;
		return -ENOSPC;
	}
	jobs[pos] = me;
	return pos;
}

//从令牌桶中奖令牌取走, 取n个, 取得是下标为td的令牌
int fetch_token (int td, int n) 
{
	int ret = 0;
	//桶的令牌文件描述符必须>0, 且去令牌的个数必须>0
	if(!((td >= 0) && n > 0)) {
		return -EINVAL;
	}
	while (jobs[td]->token <= 0) {
		//暂时没有令牌可取
		pause ();  
	}
	//当令牌桶中有令牌的时候, 才能将令牌取走 
	if (n > jobs[td]->token) {
		ret = jobs[td]->token;
	} else {
		ret = n;
	}
	jobs[td]->token -= ret;
	return ret;
}

int tbf_destroy (int td)
{
	if (!(td >= 0)) {
		return -EINVAL;
	}
	free (jobs[td]);
	jobs[td] = NULL;
	return 0;
}

//销毁令牌桶
void tbf_destroy_all (void)
{
	int i = 0;
	//停止信号模块
	sig_module_unload (); 
	for (i = 0; i < TBF_MAX; i++) {
		if (jobs[i]) {
			tbf_destroy (i);
		}
	}
	return ;
}

//找到令牌桶数组中第一个空闲位置
static int get_free_pos (void) 
{
	int i = 0;
	for (i = 0; i < TBF_MAX; i++) {
		if(jobs[i] == NULL) 
		{
			return i;
		}
	}
	return -1;
}

/*信号模块: 周期性的产生 sigalarm alarm(2)*/
void  sig_module_load (void)
{
	alarm_save = signal (SIGALRM, sig_handler);
	alarm(1);
}

void sig_handler (int s)
{
	int i = 0;
	alarm (1);
	//每当信号到来的时候, 遍历令牌桶数组 
	for (i = 0; i < TBF_MAX; i++) {
		if (jobs[i]) {
			//积攒令牌
			jobs[i]->token += jobs[i]->cps;
			if (jobs[i]->token >= jobs[i]->burst) {
				jobs[i]->token = jobs[i]->burst;
			}
		}
	}
	return ;
}
//恢复信号默认的行为 
void sig_module_unload (void)
{
	signal (SIGALRM, alarm_save);
	alarm (0);
}
